import mplhep
import matplotlib.pyplot as plt
plt.style.use(mplhep.style.LHCb2)

import pandas as pd
import numpy as np

def plot_time_distribution(file_path, save_path, bp_time = -1., bin_num = 100):

    fig = plt.figure(figsize=[12,6])

    plt.clf()
    #plt.ylim(-90., 90.)

    df = pd.read_csv(file_path)
    time_array = np.array(df['Time'])

    # 计算SOM直方图
    hist, bins = np.histogram(time_array, bin_num)
    errors = np.sqrt(hist)
    bin_width = bins[1] - bins[0]
    bin_centers = (bins[:-1] + bins[1:]) / 2

    plt.ylim(0, np.max(hist)*1.2)
    # 画BP线
    plt.plot(np.ones(2)*bp_time, np.linspace(0, np.max(hist)*0.8, 2), 'r--', label="BP")
    # 画SOM直方图
    plt.errorbar(bin_centers, hist, xerr=bin_width/2., yerr=errors, fmt='none', ecolor='black', capsize=1., elinewidth=1.5, label="SOM")
    plt.legend(fontsize='x-large')

    ax = fig.get_axes()[0]
    ax.set_xlabel("Task time / day")
    ax.set_ylabel("Candidates")

    plt.savefig(save_path)

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Test Operation module")
    #parser.add_argument("-f", "--file", default="result.csv", dest="file")
    parser.add_argument(dest="file", type=str)
    parser.add_argument("-s", "--save_path", default="default", dest="save_path")
    parser.add_argument("-t", "--bp_time", type=float, default=-1., dest="bp_time")
    arglist = parser.parse_args()
    
    import os
    csv_file = arglist.file
    if not os.path.exists(csv_file):
        raise FileNotFoundError(f"No such file or directory: '{csv_file}'.")
    save_path = csv_file[:csv_file.rfind("/")] + "/time_distribution.png"
    if arglist.save_path != "default":
        save_path = arglist.save_path

    bp_time = arglist.bp_time
    if arglist.bp_time < 0:
        bp_time = 1940.337909100867

    plot_time_distribution(file_path=arglist.file, save_path=save_path, bp_time=bp_time)

if __name__ == "__main__":
    main()